'''
这个代码占用显存较高（>5G）, 优化后的版本在matching.py
'''

from pathlib import Path
import logging
from utils.tools import *
from Detectors.superpoint.superpoint import SuperPoint

import matplotlib.pyplot as plt

class SuperPointDetector(object):
    default_config = {
        "descriptor_dim": 256,
        "nms_radius": 4,
        "keypoint_threshold": 0.005,
        "max_keypoints": -1,
        "remove_borders": 4,
        "path": Path(__file__).parent / "superpoint/superpoint_v1.pth",
        "cuda": True
    }

    def __init__(self, config={}):
        self.config = self.default_config
        self.config = {**self.config, **config}
        logging.info("SuperPoint detector config: ")
        logging.info(self.config)

        # self.device = 'cuda' if torch.cuda.is_available() and isinstance(self.config["cuda"], int)else 'cpu'
        # print(f"device: {self.device}")
        if("device" in config.keys()):
            self.device = self.config["device"]
        else:
            self.device = 'cuda' if torch.cuda.is_available() and isinstance(self.config["cuda"], int)else 'cpu'




        logging.info("creating SuperPoint detector...")
        self.superpoint = SuperPoint(self.config).to(self.device)

    def __call__(self, image):
        if len(image.shape) == 3 and image.shape[2]==3:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        logging.debug("detecting keypoints with superpoint...")

        # 占显存的是 pred 不是 image_tensor
        image_tensor = image2tensor(image, self.device)
        pred = self.superpoint({'image': image_tensor})

        # ret_dict = {
        #     "image_size": np.array([image.shape[0], image.shape[1]]),
        #     "torch": pred,
        #     "keypoints": pred["keypoints"][0].cpu().detach().numpy(),
        #     "scores": pred["scores"][0].cpu().detach().numpy(),
        #     "descriptors": pred["descriptors"][0].cpu().detach().numpy().transpose()
        # }
        #
        # return ret_dict
        pred['image'] = image_tensor

        return pred


if __name__ == "__main__":
    # img = cv2.imread("/home/daybeha/Documents/Dataset/Kitti/sequences/00/image_0/000005.png")
    img = cv2.imread("/home/daybeha/Documents/Dataset/remode_test_data/images/scene_000.png")
    # img = cv2.imread("/home/daybeha/Documents/github/DeepLabV3_ws/src/superglue/assets/scannet_sample_images/scene0711_00_frame-001680.jpg")

    detector = SuperPointDetector({"cuda": 0})
    kptdescs = detector(img)

    img = plot_keypoints(img, kptdescs["keypoints"].cpu().detach().numpy(), kptdescs["scores"].cpu().detach().numpy())
    #
    # out = 255*np.ones((H, W), np.uint8)
    # out[:H0, :W0] = image0
    # out[:H1, W0+margin:] = image1
    # out = np.stack([out]*3, -1)
    # kpts0, kpts1 = np.round(kpts0).astype(int), np.round(kpts1).astype(int)
    # white = (255, 255, 255)
    # black = (0, 0, 0)
    # for x, y in kpts0:
    #     cv2.circle(out, (x, y), 2, black, -1, lineType=cv2.LINE_AA)
    #     cv2.circle(out, (x, y), 1, white, -1, lineType=cv2.LINE_AA)
    # for x, y in kpts1:
    #     cv2.circle(out, (x + margin + W0, y), 2, black, -1,
    #                lineType=cv2.LINE_AA)
    #     cv2.circle(out, (x + margin + W0, y), 1, white, -1,
    #                lineType=cv2.LINE_AA)

    cv2.imshow("SuperPoint", img)
    cv2.waitKey()
